# Copyright (c) 2023.Huawei Technologies Co., Ltd. All rights reserved.
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at

#      http://www.apache.org/licenses/LICENSE-2.0

# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import os
import sys
import queue
import time
from datetime import datetime
import multiprocessing

import cv2
import numpy as np
import asyncio
import websockets
import pyaudio
from PyQt5.QtWidgets import *
from PyQt5.QtGui import *
from PyQt5.QtCore import QTimer, pyqtSignal, QObject, QThread

ip = 'xx.xx.xx.xx'
port1 = 'xxxx'
port2 = 'xxxx'

RECV_URL = "ws://{}:{}".format(ip, port1)
SEND_URL = "ws://{}:{}".format(ip, port2)

AUDIO_CHUNK_SIZE = 8000  # 0.5s的音频
AUDIO_SAMPLE_RATE = 16000
AUDIO_CHANNELS = 1
AUDIO_DATA_TYPE_WITDH = 2 # int16
AUDIO_FRAMES_PER_BUFFER = 1024

MAIN_UI_WIDTH = 320
MAIN_UI_HEIGHT = 500
SUBTITLES_HEIGHT = 100
RECORD_BUTTON_WIDTH = 90
RECORD_BUTTON_HEIGHT = 30
LABEL_WIDTH = 130
LABEL_HEIGHT = 15


class RecordAudioThread(QThread):
    def __init__(self, send_queue):
        super().__init__()
        self.start_record_audio = False
        self.exit_thread = False
        self.send_queue = send_queue
        self.chunk = AUDIO_CHUNK_SIZE 
        self.frames = []
        self.audio = pyaudio.PyAudio()
        self.stream = self.audio.open(format=pyaudio.paInt16, channels=AUDIO_CHANNELS,
            rate=AUDIO_SAMPLE_RATE, input=True, frames_per_buffer=self.chunk)

    def run(self):
        while not self.exit_thread:
            if self.start_record_audio:
                data = self.stream.read(self.chunk)
                self.send_queue.put(data)
        self.stream.close()

    def exitThread(self):
        self.exit_thread = True

    def startRec(self):
        self.start_record_audio = True

    def stopRec(self):
        self.start_record_audio = False


class PlayAudioThread(QThread):
    updateImage = pyqtSignal(np.ndarray)
    updateSubtitles = pyqtSignal(str)
    def __init__(self, recv_queue, subtitles_queue):
        super().__init__()
        self.exit_thread = False
        self.recv_queue = recv_queue
        self.subtitles_queue = subtitles_queue
        
        self.play = pyaudio.PyAudio()
        self.audio_stream = self.play.open(
            format=self.play.get_format_from_width(AUDIO_DATA_TYPE_WITDH),
            channels=AUDIO_CHANNELS,
            rate=AUDIO_SAMPLE_RATE,
            output=True,
            frames_per_buffer=AUDIO_FRAMES_PER_BUFFER
            )

    def exitThread(self):
        self.exit_thread = True
        self.audio_stream.close()
        self.play.terminate()

    def run(self):
        while not self.exit_thread:
            img_data, audio_data = self.recv_queue.get() # 无数据的时候阻塞
            if not self.subtitles_queue.empty():
                subtitles = self.subtitles_queue.get()
                self.updateSubtitles.emit(subtitles)
            img_buffer_numpy = np.frombuffer(img_data, dtype=np.uint8) 
            img = cv2.imdecode(img_buffer_numpy, cv2.IMREAD_COLOR)
            img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
            self.updateImage.emit(img)
            start_time = time.time()
            self.audio_stream.write(audio_data) # 这里写入的是Byte数据,不是numpy array, 阻塞式.
            end_time = time.time()
            cost_time = end_time - start_time


class Ui_MainWindow(QWidget):
    startRecordAudioSignal = pyqtSignal()
    stopRecordAudioSignal = pyqtSignal()
    stopSignal = pyqtSignal()

    def __init__(self, recv_queue, subtitles_queue, send_queue, recv_pro, send_pro, parent=None):
        super(Ui_MainWindow, self).__init__(parent)
        self.ui_init()

        self.recv_queue = recv_queue
        self.subtitles_queue = subtitles_queue
        self.send_queue = send_queue

        self.recv_pro = recv_pro
        self.send_pro = send_pro

        self.recordAudioThread = RecordAudioThread(self.send_queue)
        self.startRecordAudioSignal.connect(self.recordAudioThread.startRec)
        self.stopRecordAudioSignal.connect(self.recordAudioThread.stopRec)
        self.stopSignal.connect(self.recordAudioThread.exitThread)
        self.recordAudioThread.start()

        self.playAudioThread = PlayAudioThread(self.recv_queue, self.subtitles_queue)
        self.playAudioThread.updateImage.connect(self.UpdateUI)
        self.playAudioThread.updateSubtitles.connect(self.UpdateSubtitles)

        self.stopSignal.connect(self.playAudioThread.exitThread)
        self.playAudioThread.start()

    def closeEvent(self, a0: QCloseEvent) -> None:
        self.stopSignal.emit()
        self.recv_pro.terminate()
        self.send_pro.terminate()

    def ui_init(self):
        self.setWindowTitle("昇腾数字人参考样例")
        self.resize(MAIN_UI_WIDTH, MAIN_UI_HEIGHT)
        self.main_layout = QVBoxLayout()
        self.label_image = QLabel()
        image = cv2.imread(os.path.join('imgs', 'snips.png'))
        self.static_image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
        self.show_image = QImage(self.static_image.data,
            self.static_image.shape[1],
            self.static_image.shape[0], QImage.Format_RGB888)
        self.label_image.setPixmap(QPixmap.fromImage(self.show_image))

        self.subtitles = QListWidget(self) 
        self.subtitles.resize(MAIN_UI_WIDTH, SUBTITLES_HEIGHT)

        self.labelRec = QLabel('')
        self.labelRec.setFixedSize(LABEL_WIDTH, LABEL_HEIGHT)
        self.hbox = QHBoxLayout()
        self.recbtn = QPushButton('开始录音')
        self.recbtn.setFixedSize(RECORD_BUTTON_WIDTH, RECORD_BUTTON_HEIGHT)
        self.recbtn.clicked.connect(self.startRecordAudio)
        self.stopbtn = QPushButton('停止录音')
        self.stopbtn.setEnabled(False)
        self.stopbtn.setFixedSize(RECORD_BUTTON_WIDTH, RECORD_BUTTON_HEIGHT)
        self.stopbtn.clicked.connect(self.stopRecordAudio)
        self.hbox.addWidget(self.recbtn)
        self.hbox.addWidget(self.stopbtn)
        self.main_layout.addWidget(self.label_image)
        self.main_layout.addWidget(self.subtitles)
        self.main_layout.addWidget(self.labelRec)
        self.main_layout.addLayout(self.hbox)
        self.setLayout(self.main_layout)

    def startRecordAudio(self):
        self.labelRec.setText('◉ 正在录音...')
        self.recbtn.setEnabled(False)
        self.stopbtn.setEnabled(True)
        self.startRecordAudioSignal.emit()
        
    def stopRecordAudio(self):
        self.labelRec.setText('录音停止')
        self.recbtn.setEnabled(True)
        self.stopbtn.setEnabled(False)
        self.stopRecordAudioSignal.emit()

    def UpdateUI(self, opencv_image):
        show_image = QImage(opencv_image.data,
            opencv_image.shape[1],
            opencv_image.shape[0], QImage.Format_RGB888)
        self.label_image.setPixmap(QPixmap.fromImage(show_image))
        self.repaint()

    def UpdateSubtitles(self, sentence):
        item = QListWidgetItem(sentence)
        self.subtitles.addItem(item)
        self.subtitles.setCurrentRow(self.subtitles.count() - 1)


async def send(websocket, send_queue):
    while True:
        try:
            audio_data = send_queue.get()
            await websocket.send(audio_data)
            echo = await websocket.recv()
        except Exception as e:
            print('websocket send Exception {}'.format(e))
            break
            
async def receive(websocket, recv_queue, subtitles_queue):
    while True:
        try:
            sentence = await websocket.recv()
            subtitles_queue.put(sentence)
            if sentence:
                image_num_str = await websocket.recv()
                image_num = int(image_num_str)
                for i in range(image_num):
                    img_data = await websocket.recv() 
                    audio_data = await websocket.recv()
                    recv_queue.put((img_data, audio_data))
        except Exception as e:
            print('websocket receive Exception {}'.format(e))
            break

async def client_recv(recv_queue, subtitles_queue):
    try:
        async with websockets.connect(RECV_URL, ping_interval=None) as websocket:
            b = asyncio.get_event_loop().create_task(receive(websocket, recv_queue, subtitles_queue))
            await b
            print("client_recv wait")
    except Exception as e:
        print('client_recv receive Exception {}'.format(e))


async def client_send(send_queue):
    try:
        async with websockets.connect(SEND_URL, ping_interval=None) as websocket:
            a = asyncio.get_event_loop().create_task(send(websocket, send_queue))
            await a
            print("client_send wait")
    except Exception as e:
        print('client_send receive Exception {}'.format(e))


def recv_process(recv_queue, subtitles_queue):
    try:
        asyncio.get_event_loop().run_until_complete(client_recv(recv_queue, subtitles_queue))
        asyncio.get_event_loop().run_forever()
    except Exception as e:
        print('recv_process receive Exception {}'.format(e))


def send_process(send_queue):
    try:
        asyncio.get_event_loop().run_until_complete(client_send(send_queue))
        asyncio.get_event_loop().run_forever()
    except Exception as e:
        print('send_process receive Exception {}'.format(e))


if __name__ == '__main__':
    recv_queue = multiprocessing.Queue()
    subtitles_queue = multiprocessing.Queue()
    send_queue = multiprocessing.Queue()

    recv_pro = multiprocessing.Process(target=recv_process,
        args=(recv_queue, subtitles_queue, ))
    recv_pro.start()

    send_pro = multiprocessing.Process(target=send_process,
        args=(send_queue, ))
    send_pro.start()

    app = QApplication(sys.argv)
    ui = Ui_MainWindow(recv_queue, subtitles_queue, send_queue, recv_pro, send_pro)
    ui.show()
    sys.exit(app.exec_())
